import numpy as np
import numba as nb
from numba import cuda
import math
from ..utils import get_grid_1d, flatten


@cuda.jit()
def _gpu_cos(a: np.ndarray, b: np.ndarray):
    x = cuda.grid(1)
    if x < a.shape[0]:
        b[x] = math.cos(a[x])


def cos(a: np.ndarray):
    assert isinstance(a, np.ndarray)

    b = np.empty_like(a)

    grid, block = get_grid_1d(a.size)
    _gpu_cos[grid, block](flatten(a), flatten(b))

    return b
